from flask import Flask, Response
from flask_cors import CORS
import cv2 as cv
import time
import threading
import queue
import logging

app = Flask(__name__)
CORS(app)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CameraManager:
    def __init__(self):
        self.camera = None
        self.frame_queue = queue.Queue(maxsize=2)  # 限制队列大小，避免内存堆积
        self.capture_thread = None
        self.running = False
        self.clients = set()  # 追踪活跃客户端
        self.lock = threading.Lock()
        
    def init_camera(self):
        """初始化摄像头"""
        try:
            self.camera = cv.VideoCapture(0)
            
            if not self.camera.isOpened():
                logger.error("无法打开摄像头")
                return False
                
            # 设置摄像头参数
            self.camera.set(cv.CAP_PROP_FRAME_WIDTH, 640)
            self.camera.set(cv.CAP_PROP_FRAME_HEIGHT, 480)
            self.camera.set(cv.CAP_PROP_FPS, 30)
            self.camera.set(cv.CAP_PROP_BUFFERSIZE, 1)  # 减少缓冲区大小
            
            # 获取实际设置的参数
            width = int(self.camera.get(cv.CAP_PROP_FRAME_WIDTH))
            height = int(self.camera.get(cv.CAP_PROP_FRAME_HEIGHT))
            fps = self.camera.get(cv.CAP_PROP_FPS)
            
            logger.info(f"摄像头初始化成功 - 分辨率: {width}x{height}, 帧率: {fps}")
            return True
            
        except Exception as e:
            logger.error(f"摄像头初始化异常: {e}")
            return False
    
    def start_capture(self):
        """启动摄像头捕获线程"""
        if self.capture_thread and self.capture_thread.is_alive():
            return True
            
        if not self.init_camera():
            return False
            
        self.running = True
        self.capture_thread = threading.Thread(target=self._capture_frames)
        self.capture_thread.daemon = True
        self.capture_thread.start()
        logger.info("摄像头捕获线程已启动")
        return True
    
    def stop_capture(self):
        """停止摄像头捕获"""
        self.running = False
        if self.capture_thread and self.capture_thread.is_alive():
            self.capture_thread.join(timeout=2)
        if self.camera:
            self.camera.release()
            self.camera = None
        logger.info("摄像头捕获已停止")
    
    def _capture_frames(self):
        """摄像头捕获线程函数"""
        target_fps = 10
        frame_interval = 1.0 / target_fps
        last_capture_time = 0
        
        while self.running:
            try:
                current_time = time.time()
                
                # 控制帧率
                if current_time - last_capture_time < frame_interval:
                    time.sleep(0.01)  # 短暂休眠
                    continue
                
                if not self.camera or not self.camera.isOpened():
                    logger.warning("摄像头连接丢失，尝试重连...")
                    if not self.init_camera():
                        time.sleep(1)
                        continue
                
                success, frame = self.camera.read()
                if not success:
                    logger.warning("读取帧失败")
                    time.sleep(0.1)
                    continue
                
                # 编码为JPEG
                encode_param = [int(cv.IMWRITE_JPEG_QUALITY), 80]
                ret, buffer = cv.imencode('.jpg', frame, encode_param)
                
                if ret:
                    # 清空队列中的旧帧，只保留最新帧
                    while not self.frame_queue.empty():
                        try:
                            self.frame_queue.get_nowait()
                        except queue.Empty:
                            break
                    
                    # 放入新帧
                    try:
                        self.frame_queue.put_nowait(buffer.tobytes())
                        last_capture_time = current_time
                    except queue.Full:
                        pass  # 队列满了就跳过
                        
            except Exception as e:
                logger.error(f"捕获帧时发生错误: {e}")
                time.sleep(0.1)
    
    def add_client(self, client_id):
        """添加客户端"""
        with self.lock:
            self.clients.add(client_id)
            logger.info(f"客户端 {client_id} 已连接，当前客户端数: {len(self.clients)}")
            
            # 如果是第一个客户端，启动摄像头
            if len(self.clients) == 1:
                if not self.start_capture():
                    self.clients.discard(client_id)
                    return False
            return True
    
    def remove_client(self, client_id):
        """移除客户端"""
        with self.lock:
            self.clients.discard(client_id)
            logger.info(f"客户端 {client_id} 已断开，当前客户端数: {len(self.clients)}")
            
            # 如果没有客户端了，停止摄像头
            if len(self.clients) == 0:
                self.stop_capture()
    
    def get_frame(self):
        """获取最新帧"""
        try:
            return self.frame_queue.get(timeout=1.0)
        except queue.Empty:
            return None
    
    def get_status(self):
        """获取摄像头状态"""
        camera_info = {}
        if self.camera and self.camera.isOpened():
            camera_info = {
                "width": int(self.camera.get(cv.CAP_PROP_FRAME_WIDTH)),
                "height": int(self.camera.get(cv.CAP_PROP_FRAME_HEIGHT)),
                "fps": self.camera.get(cv.CAP_PROP_FPS)
            }
        
        return {
            "status": "running" if self.running else "stopped",
            "active_clients": len(self.clients),
            "camera_initialized": self.camera is not None and self.camera.isOpened(),
            "camera_info": camera_info
        }

# 全局摄像头管理器
camera_manager = CameraManager()

def generate_frames(client_id):
    """为单个客户端生成帧流"""
    logger.info(f"开始为客户端 {client_id} 生成帧流")
    
    try:
        while True:
            frame_data = camera_manager.get_frame()
            if frame_data is None:
                # 没有帧数据，可能摄像头出问题了
                logger.warning(f"客户端 {client_id} 没有获取到帧数据")
                time.sleep(0.1)
                continue
            
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_data + b'\r\n')
                   
    except GeneratorExit:
        logger.info(f"客户端 {client_id} 流生成器正常退出")
    except Exception as e:
        logger.error(f"客户端 {client_id} 流生成错误: {e}")
    finally:
        camera_manager.remove_client(client_id)

@app.route('/stream')
def video_stream():
    """视频流API"""
    client_id = f"client_{int(time.time() * 1000)}"
    
    if not camera_manager.add_client(client_id):
        return "摄像头初始化失败", 500
    
    return Response(generate_frames(client_id),
                    mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/status')
def status():
    """服务状态API"""
    return camera_manager.get_status()

@app.route('/camera/resolution/<int:width>/<int:height>')
def set_resolution(width, height):
    """动态调整摄像头分辨率"""
    if camera_manager.camera and camera_manager.camera.isOpened():
        try:
            camera_manager.camera.set(cv.CAP_PROP_FRAME_WIDTH, width)
            camera_manager.camera.set(cv.CAP_PROP_FRAME_HEIGHT, height)
            
            actual_width = int(camera_manager.camera.get(cv.CAP_PROP_FRAME_WIDTH))
            actual_height = int(camera_manager.camera.get(cv.CAP_PROP_FRAME_HEIGHT))
            
            return {
                "success": True,
                "requested": {"width": width, "height": height},
                "actual": {"width": actual_width, "height": actual_height}
            }
        except Exception as e:
            return {"success": False, "error": str(e)}
    else:
        return {"success": False, "error": "摄像头未初始化"}

if __name__ == '__main__':
    logger.info("正在启动视频流服务器...")
    logger.info("直播地址: http://0.0.0.0:5111/stream")
    logger.info("状态接口: http://0.0.0.0:5111/status")
    logger.info("分辨率设置: http://0.0.0.0:5111/camera/resolution/宽度/高度")
    
    try:
        app.run(host='0.0.0.0', port=5111, debug=False, threaded=True)
    except KeyboardInterrupt:
        logger.info("正在关闭服务器...")
    finally:
        camera_manager.stop_capture()
        logger.info("服务器已关闭")